#!/usr/bin/env python3
"""
数据库迁移脚本
用于向file_uploads表添加新的分析状态列
"""
import os
import sys
import sqlite3

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入项目配置
try:
    from app.core.database import DATABASE_URL
    # 从数据库URL中提取SQLite文件路径
    db_path = DATABASE_URL.replace("sqlite:///", "")
except:
    # 如果导入失败，使用默认路径
    db_path = "./flowscape.db"

def migrate_database():
    """执行数据库迁移，添加新的列"""
    print(f"开始数据库迁移...")
    print(f"数据库文件: {db_path}")
    
    try:
        # 连接到SQLite数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查file_uploads表是否存在
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='file_uploads';")
        if not cursor.fetchone():
            print("错误: file_uploads表不存在！")
            return False
        
        # 检查每个列是否已经存在
        cursor.execute("PRAGMA table_info(file_uploads);")
        existing_columns = {col[1] for col in cursor.fetchall()}
        
        # 需要添加的列
        columns_to_add = [
            ("departure_retention_status", "TEXT DEFAULT '未分析' COMMENT '离场/停留识别状态: 未分析/已分析/分析中'"),
            ("background_activity_status", "TEXT DEFAULT '未分析' COMMENT '背景/活动区分状态: 未分析/已分析/分析中'"),
            ("anomaly_detection_status", "TEXT DEFAULT '未分析' COMMENT '异常检测状态: 未分析/已分析/分析中'")
        ]
        
        # 添加缺失的列
        added_columns = []
        for col_name, col_def in columns_to_add:
            if col_name not in existing_columns:
                try:
                    # 在SQLite中，ALTER TABLE ADD COLUMN语法
                    cursor.execute(f"ALTER TABLE file_uploads ADD COLUMN {col_name} {col_def.split(' COMMENT')[0]};")
                    added_columns.append(col_name)
                    print(f"✓ 添加列 {col_name} 成功")
                except sqlite3.OperationalError as e:
                    print(f"✗ 添加列 {col_name} 失败: {e}")
            else:
                print(f"✓ 列 {col_name} 已存在，跳过")
        
        # 提交更改
        if added_columns:
            conn.commit()
            print(f"\n数据库迁移完成！成功添加了 {len(added_columns)} 个列")
        else:
            print("\n数据库已经是最新的，无需迁移")
        
        return True
        
    except Exception as e:
        print(f"数据库迁移失败: {e}")
        return False
    finally:
        # 关闭数据库连接
        if 'conn' in locals():
            conn.close()

def create_table_if_not_exists():
    """如果表不存在，则创建完整的表结构"""
    print("\n检查是否需要创建表...")
    
    try:
        # 导入SQLAlchemy相关模块
        from app.core.database import engine
        from app.models.upload import Base
        
        # 使用SQLAlchemy创建表（如果不存在）
        Base.metadata.create_all(bind=engine)
        print("✓ 表结构检查完成")
        return True
        
    except Exception as e:
        print(f"创建表失败: {e}")
        return False

def main():
    """主函数"""
    print("=" * 60)
    print("数据库迁移工具")
    print("=" * 60)
    
    # 执行数据库迁移
    migrate_database()
    
    # 确保表结构正确
    create_table_if_not_exists()
    
    print("\n" + "=" * 60)
    print("迁移工具执行完成！")
    print("您现在可以运行 insert_example.py 来测试数据插入")
    print("=" * 60)

if __name__ == "__main__":
    main()